/*	Interactive Gravitational Simulator
	Project: IGS Core
	File: NBodySystem.cpp
	Author: Mike Bantegui <mbante2@gmail.com>, Hofstra University 
	Copyright (C) 2012 - 2013 Mike Bantegui

	This file is part of the IGS software package for simulating
	the N-Body problem in real-time. This was developed as part
	of a departmental honors thesis project at Hofstra University
	in the Spring 2012 semester.

	IGS is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at your option) any later version.

	IGS is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

#include "Body.h"
#include "NBodySystem.h"
#include "Snapshot.h"
#include "ViewNode.h"
#include <fstream>

namespace IGS
{

double GetMass(Body *bodies, const int n)
{
	double mass = 0;
	
	for (int i = 0; i < n; i++)
		mass += bodies[i].Mass;

	return mass;
}

// Constructs an NBodySystem of size n
NBodySystem::NBodySystem(Body *bodies, const int n) : Bodies(bodies), BodyPointers(n), N(n),
	_Mass(GetMass(bodies, n)), _Softening(0), _SofteningSquared(0), _UpdateEnergy(true),
	_UpdateBarycenter(true), _UpdateBounds(true), _UpdateMaxima(true), _Tree(NULL)
{
	double maxMass = 0;
	for (int i = 0; i < n; i++)
	{
		BodyPointers[i] = &Bodies[i];
		maxMass = std::max(maxMass, Bodies[i].Mass);
	}

	_Dynamics.MaxMass = maxMass;
}

// Destructs the NBodySystem
NBodySystem::~NBodySystem()
{
	delete[] Bodies;
}

// Updates the kinetic and potential energies for the system
void NBodySystem::UpdateEnergy()
{
	if (!_UpdateEnergy)
		return;

	_UpdateEnergy = false;
	double kinetic = 0;
	double potential = 0;

	#pragma omp parallel for reduction(+ : kinetic, potential)
	for (int i = 0; i < N; i++)
	{
		kinetic += Bodies[i].Mass * LengthSq(Bodies[i].Velocity);
		potential += Bodies[i].Mass * Bodies[i].Potential;
	}

	_Dynamics.Kinetic = 0.5 * kinetic;
	_Dynamics.Potential = 0.5 * potential;
}

// Updates the barycenter for the system
void NBodySystem::UpdateBarycenter()
{
	if (!_UpdateBarycenter)
		return;

	_UpdateBarycenter = false;
	Vector centerMass(0, 0, 0);
	Vector centerMomentum(0, 0, 0);
	Vector angularMomentum(0, 0, 0);
	Vector force(0, 0, 0);

	Vector Ir(0, 0, 0); // Principal components of the moment of inertia tensor
	Vector Io(0, 0, 0); // Off center components of the moment of inertia tensor

	for (int i = 0; i < N; i++)
	{
		const Body &body =  Bodies[i];
		const double &m = body.Mass;
		const Vector &r = body.Position;
		const Vector &v = body.Velocity;
		const Vector &a = body.Acceleration;

		centerMass += m * r;
		centerMomentum += m * v;
		angularMomentum += Cross(r, m * v);
		force += m * a;

		const double x2 = r.x * r.x;
		const double y2 = r.y * r.y;
		const double z2 = r.z * r.z;

		Ir += m * Vector(y2 + z2, x2 + z2, x2 + y2);
		Io -= m * Vector(r.x * r.y, r.x * r.z, r.y * r.z);
	}

	angularMomentum += Cross(centerMass, centerMomentum);

	// The three columns of the moment of inertia tensor
	const Vector &I1 = Vector(Ir.x, Io.x, Io.y);
	const Vector &I2 = Vector(Io.x, Ir.y, Io.z);
	const Vector &I3 = Vector(Io.y, Io.z, Ir.z);

	const Vector a = Cross(I2, I3);
	const Vector b = Cross(I3, I1);
	const Vector c = Cross(I1, I2);
	const double det = Dot(I1, a);

	// Rows of the inverse moment of inertia tensor
	const Vector invI1 = a / det;
	const Vector invI2 = b / det;
	const Vector invI3 = c / det;

	_Dynamics.CenterOfMass = centerMass / _Mass;
	_Dynamics.CenterOfMomentum = centerMomentum / _Mass;
	_Dynamics.AngularMomentum = angularMomentum;
	_Dynamics.AngularVelocity = Vector(Dot(invI1, angularMomentum),
							  Dot(invI2, angularMomentum),
							  Dot(invI3, angularMomentum));
	_Dynamics.TotalForce = force;
}

// Updates the minimum and maximum boundaries of the system
void NBodySystem::UpdateBounds()
{
	if (!_UpdateBounds)
		return;

	_UpdateBounds = false;
	Vector minB = Bodies[0].Position;
	Vector maxB = minB;
	
	for (int i = 1; i < N; i++)
	{
		const Vector &p = Bodies[i].Position;

		minB.x = std::min(minB.x, p.x);
		minB.y = std::min(minB.y, p.y);
		minB.z = std::min(minB.z, p.z);

		maxB.x = std::max(maxB.x, p.x);
		maxB.y = std::max(maxB.y, p.y);
		maxB.z = std::max(maxB.z, p.z);
	}

	_Dynamics.MinBound = minB;
	_Dynamics.MaxBound = maxB;
}

// Updates the maximal quantities of the system
void NBodySystem::UpdateMaxima()
{
	if (!_UpdateMaxima)
		return;

	_UpdateMaxima = false;
	double maxDistance = 0;
	double maxVelocity = 0;
	double maxMomentum = 0;

	for (int i = 0; i < N; i++)
	{
		const Body &body = Bodies[i];
		const double &m = body.Mass;
		const double &r = Length(body.Position);
		const double &v = Length(body.Velocity);

		maxDistance = std::max(r, maxDistance);
		maxVelocity = std::max(v, maxVelocity);
		maxMomentum = std::max(m * v, maxMomentum);
	}

	_Dynamics.MaxDistance = maxDistance;
	_Dynamics.MaxVelocity = maxVelocity;
	_Dynamics.MaxMomentum = maxMomentum;
}

// Updates all the system's dynamical quantities
void NBodySystem::Update()
{
	UpdateBounds();
	UpdateBarycenter();
	UpdateEnergy();
	UpdateMaxima();
}

// Requires the system to update it's energy
void NBodySystem::RequireUpdate()
{
	_UpdateEnergy = true;
	_UpdateBarycenter = true;
	_UpdateBounds = true;
	_UpdateMaxima = true;
}

// Places the system in the center of momentum frame
void NBodySystem::ZeroBarycenter()
{
	const Vector &cmass = CenterOfMass();
	const Vector &cmomentum = CenterOfMomentum();

	#pragma omp parallel for
	for (int i = 0; i < N; i++)
	{
		Bodies[i].Position -= cmass;
		Bodies[i].Velocity -= cmomentum;
	}
}

// Attaches the tree to the NBodySystem for use in Snapshot
void NBodySystem::AttachTree(std::shared_ptr<ViewNode> tree)
{
	_Tree = tree;
}

// Gets the current Accelerator Tree for the system
std::shared_ptr<ViewNode> NBodySystem::Tree() const
{
	return _Tree;
}

// Gets the kinetic energy of the system
const double NBodySystem::KineticEnergy()
{
	UpdateEnergy();
	return _Dynamics.Kinetic;
}

// Gets the kinetic energy of the system
const double NBodySystem::KineticEnergy() const
{
	return _Dynamics.Kinetic;
}

// Get the potential energy of the system
const double NBodySystem::PotentialEnergy()
{
	UpdateEnergy();
	return _Dynamics.Potential;
}

// Get the potential energy of the system
const double NBodySystem::PotentialEnergy() const
{
	return _Dynamics.Potential;
}

// Gets the total energy of the system
const double NBodySystem::TotalEnergy()
{
	return KineticEnergy() + PotentialEnergy();
}

// Gets the total energy of the system
const double NBodySystem::TotalEnergy() const
{
	return KineticEnergy() + PotentialEnergy();
}

// Gets the maximum distanceof a body from the origin
const double NBodySystem::MaxDistance()
{
	UpdateMaxima();
	return _Dynamics.MaxDistance;
}

// Gets the maximum distance of a body from the origin
const double NBodySystem::MaxDistance() const
{
	return _Dynamics.MaxDistance;
}

// Gets the maximum velocity of a body in the system
const double NBodySystem::MaxVelocity()
{
	UpdateMaxima();
	return _Dynamics.MaxVelocity;
}

// Gets the maximum velocity of a body in the system
const double NBodySystem::MaxVelocity() const
{
	return _Dynamics.MaxVelocity;
}

// Gets the maximum momentum of a body in the system
const double NBodySystem::MaxMomentum()
{
	UpdateMaxima();
	return _Dynamics.MaxMomentum;
}

// Gets the maximum momentum of a body in the system
const double NBodySystem::MaxMomentum() const
{
	return _Dynamics.MaxMomentum;
}

// Gets the maximum mass of a body in the system
const double NBodySystem::MaxMass() const
{
	return _Dynamics.MaxMass;
}

// Gets the center of mass of the system
const Vector NBodySystem::CenterOfMass()
{
	UpdateBarycenter();
	return _Dynamics.CenterOfMass;
}

// Gets the center of mass of the system
const Vector NBodySystem::CenterOfMass() const
{
	return _Dynamics.CenterOfMass;
}

// Gets the center of momentum of the system
const Vector NBodySystem::CenterOfMomentum()
{
	UpdateBarycenter();
	return _Dynamics.CenterOfMomentum;
}

// Gets the center of momentum of the system
const Vector NBodySystem::CenterOfMomentum() const
{
	return _Dynamics.CenterOfMomentum;
}

// Gets the angular momentum of the system
const Vector NBodySystem::AngularMomentum()
{
	UpdateBarycenter();
	return _Dynamics.AngularMomentum;
}

// Gets the angular momentum of the system
const Vector NBodySystem::AngularMomentum() const
{
	return _Dynamics.AngularMomentum;
}

// Get the angular velocity of the system
const Vector NBodySystem::AngularVelocity()
{
	UpdateBarycenter();
	return _Dynamics.AngularVelocity;
}

// Get the angular velocity of the system
const Vector NBodySystem::AngularVelocity() const
{
	return _Dynamics.AngularVelocity;
}

// Gets the total force of the system
const Vector NBodySystem::TotalForce()
{
	UpdateBarycenter();
	return _Dynamics.TotalForce;
}

// Gets the total force of the system
const Vector NBodySystem::TotalForce() const
{
	return _Dynamics.TotalForce;
}

// Gets the minimum boundary of the system
const Vector NBodySystem::MinBound()
{
	UpdateBounds();
	return _Dynamics.MinBound;
}

// Gets the minimum boundary of the system
const Vector NBodySystem::MinBound() const
{
	return _Dynamics.MinBound;
}

// Gets the maximum boundary of the system
const Vector NBodySystem::MaxBound()
{
	UpdateBounds();
	return _Dynamics.MaxBound;
}

// Gets the maximum boundary of the system
const Vector NBodySystem::MaxBound() const
{
	return _Dynamics.MaxBound;
}

// Gets the size of the system
const int NBodySystem::Size() const
{
	return N;
}

// Gets the total mass of the system
const double NBodySystem::Mass() const
{
	return _Mass;
}

// Gets the square of the softening length parameter
const double &NBodySystem::SofteningSquared() const
{
	return _SofteningSquared;
}

// Gets the softening length parameter
const double &NBodySystem::Softening() const
{
	return _Softening;
}

// Gets a refernece to the softening length parameter
void NBodySystem::SetSoftening(const double softening)
{
	_Softening = softening;
	_SofteningSquared = softening * softening;
}

// Gets an immutable reference to the i-th body in the system
const Body &NBodySystem::operator[](const int i) const
{
	return Bodies[i];
}

// Gets a mutable reference to the i-th body in the system
Body &NBodySystem::operator[](const int i)
{
	return Bodies[i];
}

// Returns a list of pointers to all the bodies
std::vector<Body *> &NBodySystem::Pointers()
{
	return BodyPointers;
}

// Takes a snapshot of the whole system at this instant
Snapshot *NBodySystem::ToSnapshot() const
{
	Snapshot *snapshot = new Snapshot(Bodies, N, _Mass, _Dynamics, _Tree);
	// Should reset the pointer, but when rendering while updating,
	// it comes out poorly
	//_Tree.reset();
	return snapshot;
}

// Reads in a NBodySystem from the specified filename.
// Returns null if there was an error
NBodySystem *NBodySystem::Read(const char *filename)
{
	std::ifstream input(filename);
	
	if (!input.is_open())
		return NULL;

	int n;
	input >> n;
	Body *bodies = new Body[n];
		
	for (int i = 0; i < n; i++)
	{
		if (!input.good())
		{
			delete[] bodies;
			return NULL;
		}

		input >> bodies[i];
	}

	input.close();

	return new NBodySystem(bodies, n);
}

// Writes the NBodySystem to the specified filename.
// Returns whether the operation was successful
bool NBodySystem::Write(const NBodySystem &system, const char *filename)
{
	std::ofstream output(filename);

	if (!output.is_open())
		return false;

	output << system.Size() << std::endl;

	for (int i = 0; i < system.Size(); i++)
		output << system[i] << std::endl;

	output.close();

	return true;
}

}
